/*
ver 1

fixt query_RT -- mixed up buffers


*/

#include <time.h>
#include <string.h>
#include "Rainbow_Table.h"
#include "generator.h"
#include "DEHT.h"
#include "misc.h"

/*calculate the size needed for unrelatedBytes in DEHT*/
int calcSizeNeeded(rainbowTablePref *rainbowHeader, Rule *rule){
	int size = 0, i = 0, length = 0;

	size += sizeof(int); /*rainbowHeader->chainLength*/
	size += ((rainbowHeader->chainLength-1) * sizeof(LONG_INDEX_PROJ)); /*seeds array size*/
	size += sizeof(int);  /*rule->numBlocks*/
	size += (sizeof(LONG_INDEX_PROJ)*rule->numBlocks); /*rule->blockRange*/
	size += sizeof(int); /*strBlockLength*/
	for (i = 0; i < rule->numBlocks; i++)
		length += strlen(rule->block[i]);
	length += rule->numBlocks-1;
	size += (length * sizeof(char));/*strBlockLength*/
	size += sizeof(int); /*rule->dict.numWords*/
	size += sizeof(LONG_INDEX_PROJ);/*rule->dict.sumWordsPowered*/
	size += (sizeof(unsigned char)*rule->dict.numWords);  /*rule->dict.wordSize*/
	length = 0;
	for (i = 0; i < rule->dict.numWords; i++)
		length += strlen(rule->dict.words[i]);
	size += length+rule->dict.numWords; /*rule->dict.words*/


	return size;
}

/*write the pre-proccesed rule+dictionary to DEHT (in unrelatedBytes part)*/
int writeRainbowDataToDEHT(DEHT *deht, rainbowTablePref *rainbowHeader, Rule *rule){
	FILE *keyFP = NULL;
	int strBlockLength = 0, i = 0;
	keyFP = deht->keyFP;

	/*go the unrelatedBytes spot in the DEHT .key file*/
	fseek(keyFP, sizeof(struct DEHTpreferences), 0);

	/*write the rainbow table pref.*/
	if (fwrite(&(rainbowHeader->chainLength), sizeof(int), 1, keyFP) != 1)
		return 0;

	if (fwrite(rainbowHeader->seedArr, sizeof(LONG_INDEX_PROJ), rainbowHeader->chainLength-1, keyFP) != rainbowHeader->chainLength-1)
		return 0;

	/*write rule data for easy processing later*/
	if (fwrite(&(rule->numBlocks), sizeof(int), 1, keyFP) != 1)
		return 0;
	if (fwrite(rule->blockRange, sizeof(LONG_INDEX_PROJ), rule->numBlocks, keyFP) != rule->numBlocks)
		return 0;
	for (i = 0; i < rule->numBlocks; i++){
		strBlockLength += strlen(rule->block[i]);
	}

	strBlockLength += (rule->numBlocks)-1;

	if (fwrite(&strBlockLength, sizeof(int), 1, keyFP) != 1)
		return 0;


	if (fwrite(rule->block[0], sizeof(char), strBlockLength, keyFP) != strBlockLength)
		return 0;

	/*write dictionary data*/
	if (fwrite(&(rule->dict.numWords), sizeof(int), 1, keyFP) != 1)
		return 0;

	if (fwrite(&(rule->dict.sumWordsPowered), sizeof(LONG_INDEX_PROJ), 1, keyFP) != 1)
		return 0;

	if (fwrite(rule->dict.wordSize, sizeof(unsigned char), rule->dict.numWords, keyFP) != rule->dict.numWords)
		return 0;

	for (i = 0; i < rule->dict.numWords; i++){
		if (fwrite(rule->dict.words[i], sizeof(unsigned char), strlen(rule->dict.words[i])+1, keyFP) != 1+strlen(rule->dict.words[i])){
			return 0;
		}
	}
	return 1;
}

/*reads from the unrelatedBytes part of DEHT the pre-processed data of rule, dictionary and seeds calculated
 * when the rainbow table was created
 * return values:
 * return 0 if fails reading or allocating mem. (use the free_PP function in that case)
 * return 1 if succeeded
 */
int readRainbowDataToDEHT(DEHT *deht, rainbowTablePref *rainbowHeader, Rule **rule){
	FILE *keyFP = NULL;
	int size = 0, i = 0, strWordsLength = 0;
	LONG_INDEX_PROJ sizeL = 0;
	char *ruleBuffer = NULL, *p = NULL;

	*rule = (Rule *)malloc(sizeof(Rule));
	if (*rule == NULL)
		return 0;

	keyFP = deht->keyFP;
	/*go the unrelatedBytes spot in the DEHT .key file*/
	fseek(keyFP, sizeof(struct DEHTpreferences), 0);

	/*read the rainbow table pref*/
	if (fread(&size, sizeof(int), 1, keyFP) != 1)
		return 0;
	rainbowHeader->chainLength = size;
	rainbowHeader->seedArr = (LONG_INDEX_PROJ *)calloc(size-1, sizeof(LONG_INDEX_PROJ));
	if (rainbowHeader->seedArr == NULL)
		return 0;
	if (fread(rainbowHeader->seedArr, sizeof(LONG_INDEX_PROJ), size-1, keyFP) != size-1){
		free(rainbowHeader->seedArr);	
		return 0;
	}
	/*read rule data*/
	/*read how many blocks the rule has*/
	if (fread(&size, sizeof(int), 1, keyFP) != 1){
		free(rainbowHeader->seedArr);
		return 0;
	}
	(*rule)->numBlocks = size;
	(*rule)->blockRange = (LONG_INDEX_PROJ *)calloc(size, sizeof(LONG_INDEX_PROJ));
	if ((*rule)->blockRange == NULL){
		free(rainbowHeader->seedArr);
		return 0;
	}
	/*read the blockRange array*/
	if (fread((*rule)->blockRange, sizeof(LONG_INDEX_PROJ), size, keyFP) != size){
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}

	(*rule)->block = (char **)calloc(size, sizeof(char *));
	if ((*rule)->block == NULL){
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	/*read how long the rule string is*/
	if (fread(&size, sizeof(int), 1, keyFP) != 1){
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	ruleBuffer = (char *)calloc(size+1, sizeof(char));
	if (ruleBuffer == NULL){
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	if (fread(ruleBuffer, sizeof(char), size, keyFP) != size){
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}

	/*set the block array*/
	p = ruleBuffer;
	for (i = 0; i < (*rule)->numBlocks; i++){
		(*rule)->block[i] = p;
		while (*p != '\0')
			p++;
		p++;
	}

	/*read dictionary data*/
	/*read number of words*/
	if (fread(&size, sizeof(int), 1, keyFP) != 1){
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	(*rule)->dict.numWords = size;
	/*read sumWordsPowered*/
	if (fread(&sizeL, sizeof(LONG_INDEX_PROJ), 1, keyFP) != 1){
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}

	(*rule)->dict.sumWordsPowered = sizeL;
	(*rule)->dict.wordSize = (unsigned char *)calloc(size, sizeof(unsigned char));
	if ((*rule)->dict.wordSize == NULL){
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	if (fread((*rule)->dict.wordSize, sizeof(unsigned char), size, keyFP) != size){
		free((*rule)->dict.wordSize);
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	for (i = 0; i < size; i++){
		strWordsLength += (*rule)->dict.wordSize[i];
	}

	(*rule)->dict.fPtr = (char *)calloc(strWordsLength+size, sizeof(char));
	if ((*rule)->dict.fPtr == NULL){
		free((*rule)->dict.wordSize);
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	(*rule)->dict.words = (char **)calloc(size, sizeof(char *));
	if ((*rule)->dict.words == NULL){
		free((*rule)->dict.fPtr);
		free((*rule)->dict.wordSize);
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}
	/*read dict. words*/
	if (fread((*rule)->dict.fPtr, sizeof(char), strWordsLength+size, keyFP) != strWordsLength+size){
		free((*rule)->dict.words);
		free((*rule)->dict.fPtr);
		free((*rule)->dict.wordSize);
		free(ruleBuffer);
		free((*rule)->block);
		free((*rule)->blockRange);
		free(rainbowHeader->seedArr);
		return 0;
	}

	p = (*rule)->dict.fPtr;
	for (i = 0; i < size; i++){
		(*rule)->dict.words[i] = p;
		while (*p != '\0')
			p++;
		p++;
	}

	return 1;
}

LONG_INDEX_PROJ *create_seedArr(int chainLength){
	int i = 0;
	LONG_INDEX_PROJ *arr = NULL;

	arr = (LONG_INDEX_PROJ *)calloc(chainLength-1, sizeof(LONG_INDEX_PROJ));
	if (arr == NULL) return NULL;

	for (i = 0; i < chainLength-1; i++){
		arr[i] = pseudo_random_generator_proj(rand());
	}
	return arr;
}

int calc_down_Chain(DEHT *deht, rainbowTablePref *rt,Rule *rule, const unsigned char *start,int steps, unsigned char *outBuffer){
	BasicHashFunctionPtr h = NULL;
	int hashLength = 0;
	unsigned char *hBuffer = NULL;
	int i = 0;
	LONG_INDEX_PROJ k = 0;
	char pswdBuffer[80] = {0};

	/*avoid calculating out of the maximum chain length*/
	if (steps >= rt->chainLength) return 0;

	parseHashFunction(deht->header.sDictionaryName, &h, &hashLength);
	if (h == NULL)
		return 0;

	hBuffer = (unsigned char *)calloc(hashLength, sizeof(unsigned char));
	if (hBuffer == NULL) return 0;

	keyCpy((unsigned char *)hBuffer, (unsigned char *)start, hashLength);
	for (i = rt->chainLength-steps-1; i < rt->chainLength-1; i++){
		k = pseudo_random_function(hBuffer, hashLength,rt->seedArr[i]);
		generator(k, rule, pswdBuffer);
		while (strcmp(pswdBuffer,"") == 0){ /* preventing the empty password in a chain -- keeps the algorithm deterministic*/
			k++;
			generator(k, rule, pswdBuffer);
		}
		cryptHash (h, pswdBuffer, hBuffer);
	}
	keyCpy((unsigned char *)outBuffer, (unsigned char *)hBuffer, hashLength);
	free(hBuffer);
	return 1;
}


/*checks if the rainbow table has the password for the hashIn given (a hash result buffer)
 * return 0 if not in the rainbow table
 * return 1 if it is, and fills pswdBuffer with the data*
 */
int query_RT(DEHT *deht, rainbowTablePref *rt,Rule *rule, unsigned char *hashIn, char *pswdBuffer){
	int i = 0, j = 0;
	unsigned char tempBuffer[80] = {0};
	char data[80] = {0};
	int hSize = 0, res = 0;
	

	LONG_INDEX_PROJ k = 0;
	BasicHashFunctionPtr h = NULL;
	parseHashFunction(deht->header.sDictionaryName, &h, &hSize);
	if (h == NULL){
		return 0;
	}

	for (j = 0; j < rt->chainLength; j++){
		calc_down_Chain(deht, rt, rule, hashIn, j, tempBuffer);
		if ((res = query_DEHT(deht, tempBuffer, hSize, (unsigned char *)data, 80)) == DEHT_STATUS_NOT_NEEDED){
			continue;
		}
		else if (res == DEHT_STATUS_FAIL){
			perror(deht->sKeyfileName);
			return 0;
		}
		else {
			for (i = 0; i < rt->chainLength-j-1; i++){
				cryptHash(h, data, (unsigned char*)tempBuffer);
				k = pseudo_random_function((unsigned char*)tempBuffer, hSize, rt->seedArr[i]);
				generator(k, rule, data);
				while (strcmp(data,"") == 0){
					k++;
					generator(k, rule, data);
				}
			}
			cryptHash(h, data, tempBuffer);
			if (keyCmp(hashIn, (unsigned char *)tempBuffer, hSize) == 0){
				strcpy(pswdBuffer, data);
				return 1;
			}
		}
	}
	return 0;
}

/*a helping function for parseParam*/
int dispatcher(char *buf, params *par){
	char *p = NULL, *q = NULL;;

	p = buf;
	while ((*p != '\0') && (isspace(*p))){
		p++;
	}
	q = p;
	while ((*q != '\0') && (!isspace(*q))){
		if (*q == '=') break;
		q++;
	}
	if (isspace(*q)) {
		*q = '\0';
		q++;
		while ((*q != '\0') && (*q != '=')) q++;
	}
	if (*q != '=') return 0;
	*q = '\0';
	q++;
	if (strcmp(p, "dictionary_name") == 0){
		p = q;
		while ((*p != '\0') && (isspace(*p))) p++;
		q = p;
		while ((*q != '\0') && (!isspace(*q))) q++;
		if (*q != '\0'){
			while ((*q != '\0') && (isspace(*q))) q++;
		}
		if (*q != '\0') return 0;
		strcpy(par->dictFile, p);

		return 1;
	}
	else if (strcmp(p, "rainbow_chain_length") == 0){
		p = q;
		while ((*p != '\0') && (isspace(*p))) p++;
		q = p;
		while ((*q != '\0') && (isdigit(*q))) q++;
		if (*q != '\0'){
			while ((*q != '\0') && (isspace(*q))) q++;
		}
		if (*q != '\0') return 0;
		par->chainLength = atoi(p);		
		return 1;
	}
	else if (strcmp(p, "entries_in_hash_table") == 0){
		p = q;
		while ((*p != '\0') && (isspace(*p))) p++;
		q = p;
		while ((*q != '\0') && (isdigit(*q))) q++;
		if (*q != '\0'){
			while ((*q != '\0') && (isspace(*q))) q++;
		}
		if (*q != '\0') return 0;
		par->numBuckets = atoi(p);	
		return 1;
	}
	else if (strcmp(p, "bucket_block_length") == 0){
		p = q;
		while ((*p != '\0') && (isspace(*p))) p++;
		q = p;
		while ((*q != '\0') && (isdigit(*q))) q++;
		if (*q != '\0'){
			while ((*q != '\0') && (isspace(*q))) q++;
		}
		if (*q != '\0') return 0;
		par->pairs = atoi(p);
		return 1;
	}
	else if (strcmp(p, "rule") == 0){
		p = q;
		while ((*p != '\0') && (isspace(*p))) p++;
		q = p;
		while ((*q != '\0') && (!isspace(*q))) q++;
		if (*q != '\0'){
			while ((*q != '\0') && (isspace(*q))) q++;
		}
		if (*q != '\0') return 0;
		strcpy(par->rule, p);
		if (checkRule(par->rule)==0){
			return 0;	

		}		

		return 1;
	}
	else {
		return 0;
	}
}

/*this parses the ini file for part 3*/
int parseParamFile(char *filename, params *p){
	char *buffer = NULL, *toker = NULL;
	char cp[40] = {0};

	p->dictFile[0] = '\0';
	p->chainLength = 0;
	p->numBuckets = 0;
	p->pairs = 0;
	p->rule[0] = '\0';
	
	buffer = load_file(filename);
	if (buffer == NULL)
		return 0;
	toker = strtok(buffer, "\n");
	while (toker != NULL){
		strcpy(cp, toker);
		if (dispatcher(cp, p) == 0){
			free(buffer);
			return 0;
		}
		toker = strtok(NULL, "\n");
	}
	/*if the function didn't fail until now - check if all param. are assigned correctly*/
	if ((*(p->dictFile) == '\0') || (p->chainLength == 0) || (p->numBuckets == 0) || (p->pairs == 0) || (*(p->rule) == '\0')){
		free(buffer);
		return 0;
	}
	free(buffer);
	return 1;
}

void initParams(params *p){
	char buffer[80] = {0};

	keyCpy((unsigned char *)p->dictFile, (unsigned char *)buffer, 80);
	keyCpy((unsigned char *)p->rule, (unsigned char *)buffer, 80);
	p->chainLength = 0;
	p->numBuckets = 0;
	p->pairs = 0;
}
